﻿#pragma once

#include <xy/geom/bspline_curve.h>
#include <xy/geom/bspline_surface.h>
#include <xy/geom/quat.h>
#include <xy/math/poly.h>
#include <xy/solver/linear.h>

namespace xy
{

namespace alg
{

template <class T, class Real = T> class interpolator_newton
{
    static constexpr auto cmp = math::compare{};

  public:
    explicit interpolator_newton() = default;

    /**
     * @brief 构造插值条件
     *
     * @param[in] x
     * @param[in] f
     */
    interpolator_newton(const std::vector<Real> &x, const std::vector<T> &f)
    {
        std::size_t n = f.size();
        for (std::size_t i = 0; i < n; i++)
            insert(x[i], f[i]);
    }

    /**
     * @brief 获得插值多项式
     *
     * @return poly<T, Real>
     */
    math::poly<T, Real> interpolate() const
    {
        std::size_t n = m_table.size();
        math::poly<T, Real> ret(0);
        math::poly<T, Real> tmp(0);
        ret.coefficient(0) = m_table[n - 1].back();
        tmp.coefficient(0) = T{1};
        for (int i = n - 2; i >= 0; i--)
        {
            tmp *= math::poly<T, Real>({T{-m_x[i + 1]}, T{1}});
            ret += tmp * m_table[i].back();
        }
        return ret;
    }

    /**
     * @brief 直接求值
     *
     * @param[in] t
     * @return T
     */
    T value(Real t) const
    {
        T result = T{};
        T diff = T{1};
        std::size_t n = m_table.size();
        for (int i = n - 1; i >= 0; i--)
        {
            result += diff * m_table[i].back();
            diff *= t - m_x[i];
        }
        return result;
    }

    /**
     * @brief 插入数据
     *
     * @param[in] x
     * @param[in] f
     */
    void insert(Real x, T f)
    {
        // count 记录重数
        static std::size_t count = 1;
        if (!m_x.empty() && cmp.less<Real>(x, m_x.back()))
            count++;
        else
            count = 1;

        m_x.push_back(x);
        m_f.push_back(f);

        m_table.emplace_back();
        for (std::size_t i = 1; i < count; i++)
        {
            std::size_t j = m_table.size() - i;
            m_table[j].push_back(m_table[j - 1][i - 1]);
        }

        m_table[m_table.size() - count].push_back(f / math::factorial(count - 1));
        for (std::size_t i = count; i < m_table.size(); i++)
        {
            std::size_t j = m_table.size() - i - 1;
            m_table[j].push_back((m_table[j + 1][i - 1] - m_table[j][i - 1]) / (m_x[j + i] - m_x[j]));
        }
    }

    /**
     * @brief neville_aitken 算法
     *
     * @param[in] x
     * @param[in] f
     * @param[in] t
     * @return T
     */
    static T neville_aitken(const std::vector<Real> &x, const std::vector<T> &f, Real t)
    {
        // 存放差分表
        std::size_t N = x.size();
        std::vector<std::vector<T>> table(N, std::vector<T>(N));

        // 从下向上计算
        for (int i = N - 1; i >= 0; i--)
        {
            table[i][0] = f[i];
            for (std::size_t j = 1; i + j < N; j++)
                table[i][j] = (table[i + 1][j - 1] * (t - x[i]) - table[i][j - 1] * (t - x[i + j])) / (x[i + j] - x[i]);
        }

        return table[0][N - 1];
    }

    /**
     * @brief 利用差分表递归计算
     *
     * @param[in] f
     * @param[in] n
     * @return T
     */
    static std::vector<T> recursion(const std::vector<T> &f, std::size_t n)
    {
        std::vector<std::vector<T>> table(f.size(), std::vector<T>(f.size() + n));
        std::copy(f.begin(), f.end(), table[0].begin());
        for (std::size_t i = 0; i < f.size(); i++)
        {
            for (int j = i - 1; j >= 0; j--)
                table[i - j][j] = table[i - j - 1][j + 1] - table[i - j - 1][j];
        }

        std::fill(table.back().begin(), table.back().end(), table.back().front());
        for (std::size_t i = 1; i < f.size(); i++)
        {
            for (std::size_t j = i + 1; j < i + 1 + n; j++)
                table[f.size() - i - 1][j] = table[f.size() - i][j - 1] + table[f.size() - i - 1][j - 1];
        }
        return table.front();
    }

  protected:
    std::vector<Real> m_x;               // 自变量
    std::vector<T> m_f;                  // 因变量
    std::vector<std::vector<T>> m_table; // 差分表
};

template <std::size_t N, class Real = double> struct interpolator_cubic_bspline
{
    using vec_t = geom::vec<N, Real>; // P

    enum KnotType
    {
        Parameter, // 使用参数值作为节点
        Average,   // 使用参数平均值作为节点
    };

    enum BoundaryCondition
    {
        Complete,
        Natural
    };

    /**
     * @brief 三次 B 样条插值：允许指定边界条件，因此比一般的插值多 2 个控制点
     * 对于 Complete 边界条件，如果只给定方向，则模长最好取总弦长
     *
     * @param[in] f 数据点
     * @param[in] lbc 左边界类型
     * @param[in] rbc 右边界类型
     * @param[in] ba
     * @param[in] bb
     * @return std::shared_ptr<geom::bspline_curve<N>>
     */
    static std::shared_ptr<geom::bspline_curve<N>> interpolate(std::vector<vec_t> f, BoundaryCondition lbc = Natural,
                                                               BoundaryCondition rbc = Natural, vec_t ba = vec_t{},
                                                               vec_t bb = vec_t{})
    {
        auto params = geom::chord_param(f);

        std::size_t p = 3;
        std::size_t K = f.size() - 1;
        std::vector<double> knots(K + 1 + p * 2);
        for (std::size_t i = 0; i < p; i++)
            knots[i] = params[0];
        for (std::size_t i = p; i < K + p + 1; i++)
            knots[i] = params[i - p];
        for (std::size_t i = K + p + 1; i < knots.size(); i++)
            knots[i] = params[K];

        math::diagmat3d tridiag(K + 1);
        for (std::size_t i = 1; i < K; i++)
        {
            auto basis = geom::basis_funcs(knots.data(), p, i + p, params[i]);
            tridiag(i, i - 1) = basis[0];
            tridiag(i, i) = basis[1];
            tridiag(i, i + 1) = basis[2];
        }

        math::matxd Q(N, K + 1);
        for (std::size_t i = 1; i < K; i++)
        {
            for (std::size_t j = 0; j < N; j++)
                Q(j, i) = f[i][j];
        }

        if (lbc == BoundaryCondition::Complete)
        {
            tridiag(0, 0) = 3 / (params[1] - params[0]);
            auto Q0 = f[0] * 3 / (params[1] - params[0]) + ba;
            for (std::size_t j = 0; j < N; j++)
                Q(j, 0) = Q0[j];
        }
        else if (lbc == BoundaryCondition::Natural)
        {
            tridiag(0, 0) = -2 * (params[1] - params[0]) - (params[2] - params[1]);
            tridiag(0, 1) = params[1] - params[0];
            double coeff = (params[2] - params[0]) * std::pow(params[1] - params[0], 2) / 6;
            auto Q0 = ba * coeff - f[0] * (params[2] - params[0]);
            for (std::size_t j = 0; j < N; j++)
                Q(j, 0) = Q0[j];
        }

        if (rbc == BoundaryCondition::Complete)
        {
            tridiag(K, K) = 3 / (params[K] - params[K - 1]);
            auto QK = f[K] * 3 / (params[K] - params[K - 1]) - bb;
            for (std::size_t j = 0; j < N; j++)
                Q(j, K) = QK[j];
        }
        else if (rbc == BoundaryCondition::Natural)
        {
            tridiag(K, K) = -2 * (params[K] - params[K - 1]) - (params[K - 1] - params[K - 2]);
            tridiag(K, K - 1) = params[K] - params[K - 1];
            double coeff = (params[K] - params[K - 2]) * std::pow(params[K] - params[K - 1], 2) / 6;
            auto QK = bb * coeff - f[K] * (params[K] - params[K - 2]);
            for (std::size_t j = 0; j < N; j++)
                Q(j, K) = QK[j];
        }

        solver::linear_thomas solver(std::move(tridiag));
        solver.perform();
        auto P = solver.solve(std::move(Q));

        std::vector<vec_t> cpts(K + 3);
        cpts[0] = f[0];
        for (std::size_t i = 0; i < K + 1; i++)
            for (std::size_t j = 0; j < N; j++)
                cpts[i + 1][j] = P(j, i);
        cpts[K + 2] = f[K];

        return std::make_shared<geom::bspline_curve<N>>(std::move(cpts), std::move(knots), 3);
    }

    /**
     * @brief 三次 B 样条插值
     *
     * @param[in] f
     * @return std::shared_ptr<geom::bspline_surface<N>>
     */
    static std::shared_ptr<geom::bspline_surface<N>> interpolate(std::vector<std::vector<vec_t>> f)
    {
        std::size_t nu = f.size();
        std::size_t nv = f.front().size();

        std::vector<double> uknots;
        std::vector<double> vknots;

        // 插值后翻转方向，便于在另一个方向插值
        // 注意每次插值以后都会多 2 个控制点
        std::vector<std::vector<vec_t>> tar2(nv + 2, std::vector<vec_t>(nu));
        for (std::size_t i = 0; i < nu; i++)
        {
            auto spline = interpolate(f[i]);
            for (std::size_t j = 0; j < nv + 2; j++)
                tar2[j][i] = spline->pole(j);

            if (i == 0)
                vknots = std::move(spline->move_knots());
        }

        std::vector<vec_t> tar3((nu + 2) * (nv + 2));
        for (std::size_t i = 0; i < nv + 2; i++)
        {
            auto spline = interpolate(tar2[i]);
            for (std::size_t j = 0; j < nu + 2; j++)
                tar3[j * (nv + 2) + i] = spline->pole(j);

            if (i == 0)
                uknots = std::move(spline->move_knots());
        }

        return std::make_shared<geom::bspline_surface<N>>(std::move(tar3), std::move(uknots), std::move(vknots), 3, 3);
    }
};

} // namespace alg

} // namespace xy
